package com.gutsyzhan.dynamicdatasourcespringbootstarter.datasource;

import com.alibaba.druid.pool.DruidDataSource;
import com.gutsyzhan.dynamicdatasourcespringbootstarter.DynamicDataSourceContextHolder;
import com.gutsyzhan.dynamicdatasourcespringbootstarter.constants.DataSourceConstant;
import com.gutsyzhan.dynamicdatasourcespringbootstarter.provider.DynamicDataSourceProvider;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.cglib.core.ReflectUtils;
import org.springframework.util.ReflectionUtils;
import org.springframework.util.StringUtils;

import javax.sql.DataSource;
import java.io.Closeable;
import java.io.IOException;
import java.lang.reflect.Method;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 实现DisposableBean接口，是因为销毁DynamicRoutingDataSource时，会回调destroy方法，我们可以在这个方法中做一些事情
 * 实现Closeable接口，是因为我们希望能将数据源关闭
 */
@Slf4j
public class DynamicRoutingDataSource extends AbstractRoutingDataSource implements InitializingBean, DisposableBean {

    private final List<DynamicDataSourceProvider> providers;

    private String primary = DataSourceConstant.PRIMARY;

    private final Map<String, DataSource> allDataSources = new ConcurrentHashMap<>();

    public DynamicRoutingDataSource(List<DynamicDataSourceProvider> providers){
        this.providers = providers;
    }

    @Override
    protected DataSource determinateDataSource() {
        String currentDataSource = DynamicDataSourceContextHolder.getCurrentDataSource();
        return getDataSource(currentDataSource);
    }

    private DataSource getDataSource(String dataSource) {
        if(StringUtils.isEmpty(dataSource)){
            //获取默认的数据源
            return getDefaultDataSource();
        }else if(allDataSources.containsKey(dataSource)){
            return allDataSources.get(dataSource);
        }
        return getDefaultDataSource();
    }

    /**
     * 返回默认的数据源
     */
    private DataSource getDefaultDataSource() {
        DataSource dataSource = allDataSources.get(primary);
        if(dataSource != null){
            return dataSource;
        }
        throw new RuntimeException("当前没有找到默认的数据源");
    }

    public void addDataSource(String ds, DataSource dataSource) {
        DataSource oldDataSource = allDataSources.put(ds, dataSource);
        if(oldDataSource != null){
            //关闭数据源
            closeDataSource(ds,oldDataSource);
        }
        log.info("数据源 {} 添加成功",ds);
    }

    private void closeDataSource(String ds, DataSource dataSource) {
        //如果是Druid数据源，由于它实现了Closeable接口，因此里面就有close方法，可以直接调用
        if(dataSource instanceof DruidDataSource dds){
            dds.close();
        }else{
            try{
                //通过反射来获取类中名为close的方法
                Method close = ReflectionUtils.findMethod(dataSource.getClass(), "close");
                if(Objects.nonNull(close)){
                    close.invoke(dataSource);
                }
            }catch (Exception e){
                log.error("关闭数据源 {} 失败，e=",ds,e);
            }
        }
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        Map<String, DataSource> dataSourceMap = new HashMap<>(16);
        providers.stream().map(DynamicDataSourceProvider::loadDataSource).forEach(dataSourceMap::putAll);
        for (Map.Entry<String, DataSource> entry : dataSourceMap.entrySet()) {
            addDataSource(entry.getKey(), entry.getValue());
        }
    }

    @Override
    public void destroy() throws Exception {
        log.info("关闭所有数据源");
        for (Map.Entry<String, DataSource> entry : allDataSources.entrySet()) {
            closeDataSource(entry.getKey(), entry.getValue());
        }
    }

    /**
     * 修改默认的数据源名称
     */
    public void setPrimary(String primary){
        this.primary = primary;
    }
}
